library(rstan)
library(car)
library(ggplot2)

source("lines styles.R")

fs <- unlist(lapply(dir(".", "*Rdata"), function(x) { substr(x, 1, nchar(x) - 6) }))

for (f in fs) {
  
  cat(f, "\n")
  
  load(paste0(f, ".Rdata"))
  if (any(names(polex) == "experience")) names(polex)[names(polex) == "experience"] <- "gamma"
  
  vars <- unlist(lapply(names(fit)[grep("^B", names(fit))], function(x) { substr(x, 2, nchar(x)) }))
  
  if (any(substr(vars, 1, 3) == "age"))
    vars <- c(vars[-grep("^age", vars)], "entryage")
  
  ## Save PDF wil some graphs
  
  cairo_pdf(sprintf("%s_plots.pdf", f), width = 10, height = 10, onefile = TRUE)
  
  titlePage("Histogram of gamma\nand scatterplots with\nunderlying variables")
  
  print(ggplot(polex, aes(x = gamma)) + geom_histogram(fill = "blue") + theme_minimal())
  
  for (v in vars) {
    
    print(ggplot(polex, aes_string(x = v, y = "gamma")) + geom_jitter(alpha = .2) +
            geom_smooth(se = FALSE, size = 2, col = "blue"))
  }
  
  titlePage("Investigation of convergence\n\nNote that Rhat should be less\nthan, say, 1.04, close to 1")
  
  print(rstan::stan_rhat(fit))
  
  par(mfrow = c(3,4))

  varNames <- names(fit)
  if (length(grep("theta", varNames)) > 0) varNames <- varNames[-grep("theta", varNames)]
  
  firstGammaIdx <- grep("gamma", varNames)[1]
  for (v in c(varNames[1:(firstGammaIdx-1)],
                         sample(varNames[firstGammaIdx:length(varNames)], 61-firstGammaIdx))) {
    
    # hist(as.matrix(fit)[, v], main = v)
    plot(as.matrix(fit)[, v], main = v, pch = 19, cex = .5)
  }
  
  par(mfrow = c(1,1))
  
  dev.off()
  
  cairo_pdf(sprintf("%s_current_leaders.pdf", f), width = 11, height = 8, onefile = TRUE)
  
  print(ggplot(polex, aes(y = gamma, x = experienceyears)) + geom_hline(yintercept = 0, col = "gray") + geom_vline(xintercept = 0, col = "gray") + geom_point(alpha = .2) + theme_minimal())
  
  if (!("gamma_low95" %in% names(polex))) {
    
    polex$gamma_low95 <- polex$gamma_high95 <- NA
    polex$gamma_low95[!is.na(polex$gamma)] <- apply(extract(fit)$gamma, 2, quantile, .025)
    polex$gamma_high95[!is.na(polex$gamma)] <- apply(extract(fit)$gamma, 2, quantile, .975)
  }
  
  d <- polex[polex$enddate >= as.Date("2017-12-31"), ]
  d <- d[!is.na(d$ruler),]
  d <- d[d$gamma < quantile(d$gamma, .12, na.rm = TRUE) | 
           d$gamma > quantile(d$gamma, .88, na.rm = TRUE) | 
           d$ruler %in% c("Trump", "Xi Jinping", "Putin", "Narendra Modi", "Angela Merkel"), ]
  
  d$label <- paste0(d$ruler, " (", d$country, ") ", d$starty)
  
  print(ggplot(d, aes(y = reorder(label, gamma), x = gamma)) + 
    geom_point() + 
    theme_minimal() +
    labs(x = "PolEx", y = "Current leaders") + 
    geom_errorbarh(aes(xmin = gamma_low95, xmax = gamma_high95), height = 0))
  
  titlePage("As a validity check,\nrange by regime.")
  
  polex$regimetype2 <- polex$dictator
  polex$regimetype2[is.na(polex$regimetype2) | polex$regimetype2 == ""] <- recode(polex$presidential, "0='Parliamentary'; 1='Presidential'; else=NA")[is.na(polex$regimetype2) | polex$regimetype2 == ""]
  polex$regimetype2 <- as.factor(polex$regimetype2)
  table(polex$regimetype2, exclude = NULL)
  
  print(ggplot(polex[!is.na(polex$regimetype2),], aes(x = reorder(regimetype2, gamma, median, na.rm = TRUE), y = gamma)) +
    geom_boxplot(fill = "lightgray") + theme_minimal() + labs(x = "Regime type", y = "PolEx"))
  
  titlePage(paste0("In case gamma is an inexperience scale,\nalso including opposite version. If\nbetas are negative, gamma represents inexperience.\nBetas: ", paste(unname(round(apply(as.matrix(fit)[, grep("^B", names(fit))], 2, mean), 2)), collapse = " ")))

  print(ggplot(d, aes(y = reorder(label, -gamma), x = -gamma)) + 
          geom_point() + 
          theme_minimal() +
          labs(x = "PolEx", y = "Current leaders") + 
          geom_errorbarh(aes(xmin = -gamma_low95, xmax = -gamma_high95), height = 0))
  
  titlePage("As a validity check,\nrange by regime.")
  
  print(ggplot(polex[!is.na(polex$regimetype2),], aes(x = reorder(regimetype2, -gamma, median, na.rm = TRUE), y = -gamma)) +
          geom_boxplot(fill = "lightgray") + theme_minimal() + labs(x = "Regime type", y = "PolEx"))
  
  dev.off()
}